import os
import streamlit as st
import time
from dotenv import load_dotenv
import random
from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import CharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import QianfanEmbeddingsEndpoint
from langchain_community.chat_models import QianfanChatEndpoint
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate


# 加载环境变量
load_dotenv()

# 从环境变量读取秘钥
os.environ["QIANFAN_ACCESS_KEY"] = os.getenv("QIANFAN_ACCESS_KEY")
os.environ["QIANFAN_SECRET_KEY"] = os.getenv("QIANFAN_SECRET_KEY")

st.set_page_config(page_title="文心·慧眼：基于竞品数据的电商差异化营销大脑", page_icon="🌌", layout="wide")

# 注入自定义 CSS 进行极大美化
def local_css():
    st.markdown("""
    <style>
        /* 全局字体与背景 */
        .stApp {
            background: linear-gradient(to bottom right, #f8f9fa, #e9ecef);
            font-family: 'PingFang SC', 'Helvetica Neue', sans-serif;
        }
        
        /* 侧边栏美化 */
        section[data-testid="stSidebar"] {
            background-color: #ffffff;
            box-shadow: 2px 0 10px rgba(0,0,0,0.05);
        }
        
        /* 标题样式 */
        h1, h2, h3 {
            color: #2c3e50;
            font-weight: 700;
        }
        h1 {
            background: -webkit-linear-gradient(45deg, #4e54c8, #8f94fb);
            -webkit-background-clip: text;
            -webkit-text-fill-color: transparent;
            padding-bottom: 10px;
        }

        /* 按钮美化 */
        div.stButton > button {
            background: linear-gradient(90deg, #4e54c8 0%, #8f94fb 100%);
            color: white;
            border: none;
            border-radius: 8px;
            padding: 0.6rem 1.2rem;
            font-weight: 600;
            transition: all 0.3s ease;
            box-shadow: 0 4px 6px rgba(50, 50, 93, 0.11), 0 1px 3px rgba(0, 0, 0, 0.08);
        }
        div.stButton > button:hover {
            transform: translateY(-2px);
            box-shadow: 0 7px 14px rgba(50, 50, 93, 0.1), 0 3px 6px rgba(0, 0, 0, 0.08);
            background: linear-gradient(90deg, #3f44a5 0%, #7d82e8 100%);
        }

        /* 消息卡片样式 */
        .chat-card {
            padding: 20px;
            border-radius: 12px;
            margin-bottom: 15px;
            box-shadow: 0 2px 8px rgba(0,0,0,0.05);
            background-color: white;
            border-left: 5px solid #4e54c8;
        }
        
        /* 耗时徽章 */
        .time-badge {
            display: inline-block;
            background-color: #e3f2fd;
            color: #1565c0;
            padding: 2px 8px;
            border-radius: 12px;
            font-size: 0.8em;
            font-weight: bold;
            margin-top: 5px;
        }
    </style>
    """, unsafe_allow_html=True)

local_css()

# ==========================================
# 2. 核心逻辑：文件处理
# ==========================================
@st.cache_resource
def process_uploaded_file(file_content, file_name):
    """读取用户上传的产品资料 -> 存入向量数据库"""
    if not os.path.exists("./temp_data"):
        os.makedirs("./temp_data")
    
    file_path = os.path.join("./temp_data", file_name)
    with open(file_path, "wb") as f:
        f.write(file_content)

    loader = TextLoader(file_path, encoding='utf-8')
    documents = loader.load()
    text_splitter = CharacterTextSplitter(chunk_size=400, chunk_overlap=50)
    docs = text_splitter.split_documents(documents)

    embeddings = QianfanEmbeddingsEndpoint()
    db = Chroma.from_documents(docs, embeddings)
    return db

def process_custom_style_files(uploaded_files):
    """处理用户上传的自定义风格参考文件"""
    combined_text = ""
    for uploaded_file in uploaded_files:
        # 读取文本内容
        content = str(uploaded_file.read(), encoding="utf-8")
        # 简单截取前1000字作为风格样本，避免Token溢出
        combined_text += f"\n【参考范文片段】：\n{content[:800]}...\n"
    return combined_text

# ==========================================
# 3. 核心逻辑：Prompt 策略 (包含新增逻辑)
# ==========================================
def get_prompt_template(style, custom_style_content=None):
    # 通用营销大脑：教 AI 怎么把参数变成卖点
    marketing_brain = """
    你是一个金牌电商营销专家。请阅读这份【产品资料】，并基于资料回答用户需求。
    
    【思维链（你要怎么思考）】：
    1. 看到"参数"，联想"痛点"。 (如：看到298g -> 联想"轻盈手感")
    2. 严格遵守资料里的"品牌调性"和"违禁词"。
    3. 🔥 **核心卖点深挖（Deep Dive）**：找出产品最核心的1-2个卖点，绝不能一笔带过！必须使用"显微镜式"描写——通过排比、场景构建、感官描述（视觉/触觉/听觉）将这部分的篇幅拉长，使其在文中占据绝对的视觉重心，体现技术应用的深度。
    """
    
    # 基础模板
    templates = {
        "小红书种草": marketing_brain + """
        【请模仿以下范文的语气】：
        1.范文一:
        标题：[产品名称]杀到！这波配置直接把我香晕了🤩
        正文：
        家人们！[产品名称]的爆料也太炸了吧！作为一个资深[品牌]粉，看完这些信息我已经坐不住了！😤
        首先说外观，[核心设计部位]要搞大事情！新增 [具体配置A] + [具体配置B]，外观走 [设计风格形容词] 的路子，辨识度直接拉满！✨
        屏幕方面更绝，[型号A] 和 [型号B] 是[xx党]狂喜的设计，不管你爱哪种，都能选到心头好～❤️
        性能这块更是王炸！全系搭载 [核心技术/系统]，搭配传说中的 [核心处理器/马达/引擎]，流畅度和功耗控制想想就激动！🚀
        [其他功能] 这些咱先不说，就这外观和性能的升级，已经让我开始疯狂期待了！[品牌名称] 这波操作，是要把 [产品品类] 的天花板再抬高一截啊！📈

        2.范文二:
        标题:对 [特定人群/来源] 的 [产品品类] 真香了！[核心卖点] 被我挖到了！
        正文：
        之前总觉得 [产品品类] 是智商税 🧠
        直到 [朋友/同事] 给我试了这支
        说是从 [可靠来源] 那安利到的
        现在 [周围圈子] 都在用！✌🏻
    
        [型号/颜色A] 上手是 [效果形容词A] ✌︎' ֊'
        另一款 [型号/颜色B] 是 [效果形容词B]
        就像 [理想状态形容词]
        [使用场景] 瞬间 [核心功效]
        
        [急救场景（如早八/加班）] 来不及 [复杂动作]
        用它直接冲 [目的地] 😎
        
        成分安心到离谱！
        [核心成分A] + [核心成分B] 不仅 [基础功效]，还能 [进阶功效] 🥰
        
        最近连 [隔壁人群] 都来问链接
        果然好东西藏不住！
        
        3.范文三:
        标题:被 [亲密关系称呼] 强行安排的 [产品名称] 初体验…结果真香了！！😭
        正文：
        来对 [产品品类] 一直是观望状态，总觉得 [刻板印象/顾虑点] 🤔
        结果某 [亲密关系称呼] 非说"[推荐理由/诱导话术]"，硬拉着我 [购买/体验] 了一套…
        
        [使用/上手] 的瞬间真的有被惊到！😳
        [核心参数/design] 意外 [体验感受/解决痛点]！！✨
        以前总觉得这种 [风格/产品] 离自己很远，没想到尝试之后居然有点上头…💘
        果然 [感叹句/总结句]！
        
        你们的 [产品名称] 初体验也是这样吗？🙈

        4.范文四:
        标题:大型真香现场！🥹我居然爱上了 [长辈称呼] 同款 [产品品类]
        正文：
        以前每次 [接触/使用] 我爸的 [旧款产品] 都偷偷吐槽 "这也太中年老干部风了吧" 🤣
        直到刷到 [新产品型号] —— 直接被狠狠戳中！💘
        这还是我认识的 [品牌/品类] 吗？！
        
        颜值暴击💥 [核心气质] 拉满！
        [外观设计A] + [外观设计B]，像 [比喻对象] 一样酷！
        再配上 [专属配置]，好几次被路人问 "这是新款 [高级品类] 吗"？
        完全打破 "[产品品类] = [刻板印象]" 的魔咒好吗！
        
        现在轮到我调侃老爸了 😎
        每次 [使用场景] 的时候，都故意在我爸面前晃 🤪
        "爸，你那老 [旧款产品] 该更新换代啦，你看我的比你还潮！"
        
        我爸嘴上不服气，却偷偷 [偷窥/体验] 了好几回我的 [产品核心部位] 😉hhh
        谁能想到啊！曾经嫌弃的 [产品品类]，现在成了我出门必炫耀的心头好 🥳
        
        果然选 [产品品类] 不能凭刻板印象，[新产品型号] 真的太懂年轻人了!😍

        5.范文五:
        标题:种草 [产品名称]！用了它之后，我的 [使用场景] 变得超 [积极形容词] 了！
        正文：
        最近入手了 [产品名称]，真的是太惊喜了！✨
        作为一个 [用户身份/背景]，我对 [产品品类] 的要求还是挺高的。
        但是这款 [产品名称] 完全超出了我的预期！😍
        首先，它的 [核心卖点A] 让我在 [使用场景] 时感到非常 [积极形容词A]。以前我总觉得 [使用场景] 是一件麻烦事，但有了它，一切都变得轻松多了！💪
        其次，它的 [核心卖点B] 也让我爱不释手。在 [使用场景B] 的时候，我能明显感觉到它带来的便利和舒适感。真的让我对 [产品品类] 有了全新的认识！🌟
        最后，我还要说的是它的 [核心卖点C]。这让我在 [使用场景C] 时，体验感提升了一个档次！简直是为我们这些 [用户身份/背景] 量身定做的啊！🎯
        总之，种草 [产品名称] 绝对没错！如果你和我一样是个 [用户身份/背景]，强烈推荐你也试试这款神器！相信我，你会爱上它的！😍🔥

        【现在轮到你了】：
        产品资料：{context}
        用户需求：{question}
        
        【生成要求 - 加强版】：
        1. 智能匹配：根据【产品资料】特性，从上方 5 个范文中选择最贴切的一个进行深度模仿。
        2. 🔥 **语气强制**：必须使用极其情绪化的"网络黑话"！
           - 强制使用词库：绝绝子、yyds、家人们、真香、拿捏、翘楚、天花板、救命、暴风吸入（如果适用）。
           - 禁止使用：虽然但是、总的来说、综上所述（太书面了）。
        3. 🔥 **Emoji 密度**：每两句话必须带一个 Emoji (✨🔥😭🥺😎)，这一条非常重要！
        4. 🛑 防幻觉指令：语气可以夸张，但**具体参数（价格、配置、型号）必须严格来自【产品资料】，严禁编造资料里没有的数据！**
        5. 结尾必须带 5 个以上的话题标签（Hashtag）。
        6. 🔥 **(新增) 核心卖点扩写**：请从产品中挑出 1 个最强痛点（例如技术参数或功效），用"痛点描述 + 夸张的解决效果 + 个人主观感受"的三层逻辑详细展开，这一段的字数要比其他段落多一倍，把它吹爆！
        """,
        
        "抖音直播脚本": marketing_brain + """
        【请模仿以下范文的语气】：
        1.范文一:
        主播：来！所有女生！看这里！👀
        主播：(展示动作) 大家看这个细节，是不是薄如蝉翼？✨
        主播：外面专柜要卖多少？5999！今天在我直播间，把价格打下来！只要2999！🔥
        主播：库存只有最后50单，3！2！1！上链接！🔗

        2.范文二:
        主播：Oh my god！所有女生！你们的魔鬼来咯！🔔
        主播：(展示特写) 你们看这个 [产品核心细节]，简直是把"高级"两个字写在了脸上！✨
        主播：老板娘在吗？原价 [原价] 我真的不敢看！今天李老头给你们炸到 [现价]！买一送一！💥
        主播：美眉们，全中国只有最后2000组，手慢无！买它！买它！买它！❗️

        3.范文三:
        主播：兄弟们！别眨眼！给我上才艺！🤜
        主播：(进行暴力测试动作) 看到没有？怎么 [破坏性动作] 都不坏！这就是质量！这就是硬货！🧱
        主播：品牌方今天不想过了是吧？[原价] 不要！给我改成 [现价]！直接上车！🚑
        主播：有没有？没抢到的扣1！后台再加500单！手速要快！给兄弟们炸！⚡️

        4.范文四:
        主播：听听！都给我听听！别去买外面那些智商税了！😒
        主播：(傲娇展示) 看看这 [产品质感]，K总什么时候骗过你们？这带出去多有面子？💅
        主播：今天这价格，我是真没赚钱，纯粹给姐姐们谋福利！[现价]！闭眼冲！🤐
        主播：只有这一波啊，抢不到别来私信我哭！去拍！链接在左下角！🛍️

        
        【现在轮到你了】：
        产品资料：{context}
        用户需求：{question}
        
        【生成要求 - 加强版】：
        1. 智能匹配：根据产品属性，选择最合适的一个主播风格。
        2. 🔥 **节奏强制**：必须是"短句"+"感叹号"！不要长篇大论！
        3. 🔥 **互动感**：必须包含与弹幕的互动，例如："没抢到的扣1"、"觉得贵的把'贵'打在公屏上"、"后台运营别睡了赶紧上库存"。
        4. 🛑 防幻觉指令：范文里的"2999"是假数据，不要抄！**请务必去【产品资料】里找真实的价格和库存！**如果资料没写价格，就说"价格炸裂"或者"给到骨折价"。
        5. ⏱️ **专业时间轴规划**：请务必按照直播带货的黄金节奏，为脚本加上明确的**粗体时间标记**。请严格参考以下结构：
            - **【0-30秒】痛点引入**：不要直接卖货！先讲生活中遇到的麻烦，引发共鸣，留住观众。
            - **【30-60秒】产品介绍**：亮出产品，展示核心卖点（对应解决痛点）。
            - **【60-90秒】高潮演示**：进行暴力测试、上脸实测或特写展示，证明效果。
            - **【90-120秒】逼单时刻**：强调价格优势、赠品和库存紧张，引导下单。
        6. 🔥 **(新增) 卖点高光时刻**：在【30-60秒】或【60-90秒】区间，对于最核心的技术/卖点，必须使用连续 3 个感叹句或反问句进行高强度强调（如："这就是X技术！懂行的兄弟们打个懂！还有谁能做到？"），增加语气的爆发力。
        """,
        
        "京东详情页": marketing_brain + """
        【请模仿以下五篇范文的语气】：
        【范文1：核心技术解析】搭载行业领先的第V代数字引擎，效率提升30%。这意味着，它能从源头上解决卡顿痛点。
        【范文2：工艺美学】机身采用航空级铝合金材质，经由CNC精密雕刻，触感细腻冷冽。
        
        【范文3：生物黑科技风（适合美妆/护肤）】
        【核心成分解析】蕴含高浓度 [核心成分名称]，深入肌底，从源头激活 [细胞/胶原] 再生。
        【夜间修护科技】搭载独家 [技术名称] 科技，精准捕捉受损因子，在 [时间段] 释放修护能量。
        
        【范文4：人体工学风（适合椅子/鞋服）】
        【仿生脊椎设计】依据 [人体工学原理]，打造 [结构名称] 支撑系统。它能自动贴合您的 [身体部位] 曲线，提供动态支撑。
        【专利网面科技】采用 [材质名称] 材质，具有独特的 [材质特性]，均匀分散身体压力。
        
        【范文5：精密机械风（适合咖啡机/电器）】
        【恒温萃取系统】配备 [加热技术名称]，精准控制水温在 [温度范围]°C 黄金萃取区间，完美释放 [原料] 的原有风味。
        【研磨工艺】内置 [磨盘材质] 锥形磨豆器，支持 [数值] 档研磨粗细调节。
        
        【现在轮到你了】：
        产品资料：{context}
        用户需求：{question}
        
        【生成要求 - 严格遵守】：
        1. 智能匹配：请根据【产品资料】的类型，自动选择上面 5 个范文中**最贴切的一种风格**进行模仿。
        2. 语气要求：专业、严谨、客观、数据支撑。多用"搭载"、"配备"、"源于"、"甄选"等高级词汇。
        3. 🛑 防幻觉指令：范文里的具体参数（如"30%"、"CNC"）仅供参考句式！**所有参数必须严格来自【产品资料】，严禁编造资料里没有的数据！**
        4. 🔥 **(新增) 技术深度解析**：针对 1-2 个核心技术参数，不能只列数据！必须采用"原理+数据+体验"的深度写法（例如：不仅说搭载了X引擎，还要解释X引擎是如何工作的，以及它让用户的使用感受提升了多少），体现专业壁垒。
        """
    }

    # 4. 🔥 【新增】自定义风格逻辑
    if style == "✨ 自定义风格 (Custom)" and custom_style_content:
        # 修复：使用字符串替换而不是 f-string 或 .format()
        custom_template = marketing_brain + """
        【你收到了一些参考资料（论文/爆文），请严格模仿其语感】：
        {custom_style_content}

        【现在轮到你了】：
        产品资料：{context}
        用户需求：{question}

        【生成要求 - 严格遵守】：
        1. 🔥 **风格迁移**：请分析上述【参考范文片段】的句式结构、用词专业度、情感色彩，并完全模仿这种风格来撰写本次的文案。
        2. 🛑 **严禁数据混淆**：参考范文中的任何具体数据（如价格、参数、型号、实验数据）**一律不准使用**！你只能使用【产品资料】里的数据。
        3. 🔥 **核心卖点深挖**：如上所述，对核心卖点进行显微镜式扩写，字数加倍。
        """
        # 直接替换占位符
        custom_template = custom_template.replace("{custom_style_content}", custom_style_content)
        return custom_template

    return templates.get(style, templates["小红书种草"])

# ==========================================
# 4. 辅助功能函数
# ==========================================

def competitor_analysis(product_context, llm):
    """🔍 竞争情报分析系统"""
    start_time = time.time() # ⏱️ 计时开始
    analysis_prompt = f"""
    基于以下产品资料，请进行简短的【竞争情报分析】：
    产品资料摘要：{product_context[:500]}...
    
    请输出以下JSON格式信息（不要输出Markdown，直接输出文本）：
    1. 【竞品对标】：列出市面上可能的2个竞争对手类型。
    2. 【差异化优势】：该产品打败竞品的1个核心杀手锏是什么？
    3. 【营销切入点】：建议从什么刁钻的角度进行宣传？
    """
    try:
        response = llm.invoke(analysis_prompt)
        end_time = time.time() # ⏱️ 计时结束
        return response.content, end_time - start_time
    except:
        return "无法连接竞争情报数据库...", 0

def virtual_review_board(generated_content, llm, style):
    """👥 虚拟评审团 (自适应风格)"""
    start_time = time.time() # ⏱️ 计时开始
    review_criteria = {
        "小红书种草": {
            "intro": "你现在是【小红书爆文鉴赏团】。专注于：颜值、种草力、情绪价值。",
            "roles": "**💅 颜值博主**：[点评美感]\n**🗣️ 互动运营**：[点评互动引导]\n**⚖️ 合规专家**：[检查违禁词]",
        },
        "抖音直播脚本": {
            "intro": "你现在是【抖音直播间操盘手】。专注于：黄金三秒、节奏感、逼单话术。",
            "roles": "**🎬 直播导演**：[点评节奏]\n**💰 投流专家**：[点评转化潜力]\n**⚖️ 合规专家**：[检查虚假宣传]",
        },
        "京东详情页": {
            "intro": "你现在是【京东电商产品经理】。专注于：专业度、参数准确性、信任背书。",
            "roles": "**👨‍🔬 产品经理**：[点评逻辑]\n**📉 转化专家**：[点评信任感]\n**⚖️ 合规专家**：[检查广告法]",
        },
        "✨ 自定义风格 (Custom)": {
            "intro": "你现在是【高级内容审核主编】。专注于：风格一致性、逻辑自洽性。",
            "roles": "**📝 主编**：[点评风格还原度]\n**🕵️ 事实核查员**：[检查是否混淆了参考资料的数据]",
        }
    }
    # 默认为小红书，防止KeyError
    current = review_criteria.get(style, review_criteria["小红书种草"])
    
    review_prompt = f"""
    {current['intro']}
    文案：{generated_content}
    
    请输出评审结果（仅输出专家会诊部分，保持客观）：
    ### 👥 专家会诊
    {current['roles']}
    
    （请针对文案进行打分和简评，不要直接重写文案）
    """
    response = llm.invoke(review_prompt)
    end_time = time.time() # ⏱️ 计时结束
    return response.content, end_time - start_time

def rewrite_content(original_content, review_feedback, llm):
    """(新增功能) 🔧 一键优化"""
    start_time = time.time() # ⏱️ 计时开始
    rewrite_prompt = f"""
    你是一个精益求精的文案大师。请根据【评审团意见】修改【原始文案】。
    
    【原始文案】：
    {original_content}
    
    【评审团意见（参考诊断）】：
    {review_feedback}
    
    【修改要求】：
    1. 必须根据诊断意见优化痛点。
    2. 保持原有的平台风格。
    3. 🔥 **重点**：对被专家表扬或指出的核心卖点，请进一步增加细节描述，强化技术应用的深度。
    4. 直接输出修改后的完整文案，不要废话。
    """
    response = llm.invoke(rewrite_prompt)
    end_time = time.time() # ⏱️ 计时结束
    return response.content, end_time - start_time

# ==========================================
# 5. 前端界面 (Streamlit)
# ==========================================
st.title("🌌 文心·慧眼：基于竞品数据的电商差异化营销大脑")
st.markdown("**🚀 融合 竞争情报分析 + RAG生成 + 虚拟评审 (v2.0 Pro)**")

# 初始化Session
if "db" not in st.session_state: 
    st.session_state.db = None
if "messages" not in st.session_state: 
    st.session_state.messages = [{"role": "assistant", "content": "👋 您好！我是您的智能营销助手。请上传产品资料，我将为您生成全案策略。"}]
if "last_review" not in st.session_state: 
    st.session_state.last_review = ""
if "raw_text" not in st.session_state:
    st.session_state.raw_text = ""
if "custom_style_data" not in st.session_state:
    st.session_state.custom_style_data = None

# --- 侧边栏 ---
with st.sidebar:
    st.image("https://cdn-icons-png.flaticon.com/512/4712/4712035.png", width=80)
    st.header("📂 资料中心")
    
    with st.expander("📝 第一步：产品资料 (必须)", expanded=True):
        uploaded_file = st.file_uploader("上传产品说明书 (TXT)", type=['txt'])
        if uploaded_file and st.session_state.db is None:
            with st.spinner("🔄 正在构建向量索引..."):
                st.session_state.db = process_uploaded_file(uploaded_file.getvalue(), uploaded_file.name)
                st.session_state.raw_text = str(uploaded_file.getvalue(), encoding='utf-8')
            st.success(f"✅ 已就绪：{uploaded_file.name}")

    st.divider()
    
    st.header("⚙️ 策略控制台")
    # 增加自定义风格选项
    style = st.selectbox("目标风格", ["小红书种草", "抖音直播脚本", "京东详情页", "✨ 自定义风格 (Custom)"])
    
    # 📂 自定义风格上传区
    if style == "✨ 自定义风格 (Custom)":
        with st.expander("📂 上传风格参考 (论文/爆文)", expanded=True):
            style_files = st.file_uploader(
                "请上传 1-10 个参考文件 (TXT/MD)", 
                type=['txt', 'md'], 
                accept_multiple_files=True
            )
            if style_files:
                if len(style_files) > 10:
                    st.warning("⚠️ 最多支持10个文件，将仅使用前10个。")
                    style_files = style_files[:10]
                st.session_state.custom_style_data = process_custom_style_files(style_files)
                st.info(f"✅ 已提取 {len(style_files)} 份文档的风格特征")
            else:
                st.session_state.custom_style_data = None
                st.warning("请上传参考文件以激活自定义风格！")

    model_name = st.selectbox("模型引擎", ["ERNIE-4.0-8K", "ERNIE-3.5-8K", "ERNIE-Speed-128K"])
    
    col_t1, col_t2 = st.columns(2)
    with col_t1:
        enable_competitor = st.toggle("🔍 竞争情报", value=True)
    with col_t2:
        enable_review = st.toggle("👥 虚拟评审", value=True)
    
    st.divider()
    if st.button("🔄 重置系统"):
        st.cache_resource.clear()
        st.session_state.clear()
        st.rerun()

# --- 主界面 ---
# 历史消息渲染
for msg in st.session_state.messages:
    # 使用自定义CSS类包裹消息
    st.markdown(f'<div class="chat-card">', unsafe_allow_html=True)
    with st.chat_message(msg["role"]):
        # 检测是否是子内容（竞争情报 或 虚拟评审），如果是则折叠显示
        if "### 🔍 竞争情报分析" in msg["content"]:
            with st.expander("🔍 竞争情报分析 (点击展开)"):
                st.markdown(msg["content"])
                # 提取耗时（如果有）
                if "cost_time" in msg:
                    st.markdown(f'<div class="time-badge">⏱️ 耗时: {msg["cost_time"]:.2f}s</div>', unsafe_allow_html=True)
        elif "### 👥 虚拟评审团" in msg["content"]:
            with st.expander("👥 虚拟评审团意见 (点击展开)"):
                st.markdown(msg["content"])
                if "cost_time" in msg:
                    st.markdown(f'<div class="time-badge">⏱️ 耗时: {msg["cost_time"]:.2f}s</div>', unsafe_allow_html=True)
        else:
            # 正常显示主文案
            st.markdown(msg["content"])
            if "cost_time" in msg:
                 st.markdown(f'<div class="time-badge">⏱️ 耗时: {msg["cost_time"]:.2f}s</div>', unsafe_allow_html=True)
    st.markdown('</div>', unsafe_allow_html=True)

# --- 输入逻辑 ---
prompt = st.chat_input("💡 输入需求（例如：帮我写个种草文案，核心卖点是续航）")

if prompt:
    # 检查自定义风格是否就绪
    if style == "✨ 自定义风格 (Custom)" and not st.session_state.custom_style_data:
        st.error("❌ 请先在左侧上传【自定义风格】参考文件！")
    elif st.session_state.db is None:
        st.error("❌ 请先在左侧上传【产品资料】文件！")
    else:
        st.session_state.messages.append({"role": "user", "content": prompt})
        # 立即渲染用户提问，不等待
        with st.chat_message("user"): 
            st.markdown(prompt)

        try:
            llm = QianfanChatEndpoint(model=model_name, temperature=0.7)
            
            # 1. 文案生成 (核心内容)
            # 传入自定义风格内容（如果有）
            current_template = get_prompt_template(style, st.session_state.custom_style_data)
            PROMPT = PromptTemplate(template=current_template, input_variables=["context", "question"])
            qa_chain = RetrievalQA.from_chain_type(
                llm=llm, 
                retriever=st.session_state.db.as_retriever(search_kwargs={"k": 2}), 
                chain_type_kwargs={"prompt": PROMPT}
            )
            
            with st.chat_message("assistant"):
                # A. 生成主文案
                start_gen = time.time() # ⏱️
                with st.spinner(f"✍️ AI ({model_name}) 正在创作..."):
                    response = qa_chain.invoke({"query": prompt})
                    generated_text = response['result']
                    gen_time = time.time() - start_gen # ⏱️
                    
                    st.markdown(f"## 📝 生成内容\n\n{generated_text}")
                    st.markdown(f'<div class="time-badge">⏱️ 耗时: {gen_time:.2f}s</div>', unsafe_allow_html=True)
                    
                    st.session_state.messages.append({
                        "role": "assistant", 
                        "content": f"## 📝 生成内容\n\n{generated_text}",
                        "cost_time": gen_time
                    })
                
                # B. 生成竞争情报
                if enable_competitor and st.session_state.raw_text:
                    with st.status("🔍 正在分析市场数据...", expanded=False) as status:
                        insight, comp_time = competitor_analysis(st.session_state.raw_text, llm)
                        st.session_state.messages.append({
                            "role": "assistant", 
                            "content": f"### 🔍 竞争情报分析\n\n{insight}",
                            "cost_time": comp_time
                        })
                        status.update(label=f"✅ 竞争情报已生成 ({comp_time:.2f}s)", state="complete")
                
                # C. 生成虚拟评审
                if enable_review:
                    with st.status("👥 专家评审团正在会诊...", expanded=False) as status:
                        review_result, review_time = virtual_review_board(generated_text, llm, style)
                        st.session_state.last_review = review_result 
                        st.session_state.messages.append({
                            "role": "assistant", 
                            "content": f"### 👥 虚拟评审团\n\n{review_result}",
                            "cost_time": review_time
                        })
                        status.update(label=f"✅ 专家评审已完成 ({review_time:.2f}s)", state="complete")
            
            st.rerun()

        except Exception as e:
            st.error(f"系统错误: {e}")

# --- 底部功能区 ---
if (st.session_state.messages and 
    st.session_state.messages[-1]["role"] == "assistant"):
    
    st.divider()
    col1, col2 = st.columns([1, 4])
    
    with col1:
        if st.session_state.get("last_review"):
            if st.button("🔧 一键优化 (Auto-Refine)", type="primary", help="根据评审团意见自动重写"):
                with st.spinner("⚡️ AI 正在根据专家意见重写文案..."):
                    llm = QianfanChatEndpoint(model=model_name, temperature=0.7)
                    
                    # 寻找最近的生成内容
                    original_msg = None
                    for i in range(len(st.session_state.messages)-1, -1, -1):
                        if "生成内容" in st.session_state.messages[i]["content"]:
                            original_msg = st.session_state.messages[i]
                            break
                    
                    if original_msg:
                        original_content = original_msg["content"].replace("## 📝 生成内容\n\n", "")
                        review = st.session_state.last_review
                        
                        new_version, rewrite_time = rewrite_content(original_content, review, llm)
                        
                        st.session_state.messages.append({
                            "role": "assistant", 
                            "content": f"## ✅ 优化版本\n\n{new_version}",
                            "cost_time": rewrite_time
                        })
                        
                        st.session_state.last_review = ""
                        st.rerun()
else:
    st.button("🔧 一键优化", disabled=True)